"""
绘图工具模块

用于可视化子通道分析的计算结果
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from .. import config as cfg

class ResultPlotter:
    """结果可视化类"""
    
    def __init__(self):
        """初始化绘图器"""
        # 设置中文显示
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        
    def create_mesh_grid(self):
        """创建网格点"""
        x = np.linspace(0, cfg.CHANNEL_LENGTH, cfg.AXIAL_NODES)
        y = np.arange(len(cfg.CHANNEL_CONNECTIONS))
        return np.meshgrid(x, y)
        
    def plot_3d_field(self, ax, X, Y, Z, title, xlabel, ylabel, zlabel):
        """
        绘制3D场分布图
        
        参数：
        ax: Axes3D, matplotlib 3D轴对象
        X, Y: array, 网格点坐标
        Z: array, 场变量值
        title: str, 图标题
        xlabel, ylabel, zlabel: str, 坐标轴标签
        """
        surf = ax.plot_surface(X, Y, Z, cmap=cm.viridis)
        ax.set_title(title)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_zlabel(zlabel)
        plt.colorbar(surf, ax=ax, shrink=0.5, aspect=5)
        ax.view_init(elev=20, azim=45)
        ax.grid(True)
        
    def plot_liquid_phase_results(self, channels):
        """
        绘制液相相关的计算结果
        
        参数：
        channels: dict, 子通道字典
        """
        # 创建图形窗口
        fig = plt.figure(figsize=(20, 15))
        
        # 创建网格点
        X, Y = self.create_mesh_grid()
        
        # 提取场变量数据
        alpha_l = np.array([channel.alpha_l for channel in channels.values()])
        u_l = np.array([channel.u_l for channel in channels.values()])
        v_l = np.array([channel.v_l for channel in channels.values()])
        T_l = np.array([channel.T_l for channel in channels.values()])
        p = np.array([channel.pressure for channel in channels.values()])
        T_w = np.array([channel.T_w for channel in channels.values()])
        
        # 1. 液相分数分布
        ax1 = fig.add_subplot(231, projection='3d')
        self.plot_3d_field(ax1, X, Y, alpha_l,
                          '液相分数分布', '轴向位置 (m)', '子通道编号', '液相分数')
        
        # 2. 液相轴向速度分布
        ax2 = fig.add_subplot(232, projection='3d')
        self.plot_3d_field(ax2, X, Y, u_l,
                          '液相轴向速度分布', '轴向位置 (m)', '子通道编号', '速度 (m/s)')
        
        # 3. 液相横向速度分布
        ax3 = fig.add_subplot(233, projection='3d')
        self.plot_3d_field(ax3, X, Y, v_l,
                          '液相横向速度分布', '轴向位置 (m)', '子通道编号', '速度 (m/s)')
        
        # 4. 液相温度分布
        ax4 = fig.add_subplot(234, projection='3d')
        self.plot_3d_field(ax4, X, Y, T_l,
                          '液相温度分布', '轴向位置 (m)', '子通道编号', '温度 (K)')
        
        # 5. 压力分布
        ax5 = fig.add_subplot(235, projection='3d')
        self.plot_3d_field(ax5, X, Y, p,
                          '压力分布', '轴向位置 (m)', '子通道编号', '压力 (Pa)')
        
        # 6. 壁面温度分布
        ax6 = fig.add_subplot(236, projection='3d')
        self.plot_3d_field(ax6, X, Y, T_w,
                          '壁面温度分布', '轴向位置 (m)', '子通道编号', '温度 (K)')
        
        plt.tight_layout()
        plt.show()
        
    def plot_vapor_phase_results(self, channels):
        """
        绘制气相相关的计算结果
        
        参数：
        channels: dict, 子通道字典
        """
        # 创建图形窗口
        fig = plt.figure(figsize=(20, 15))
        
        # 创建网格点
        X, Y = self.create_mesh_grid()
        
        # 提取场变量数据
        alpha_v = np.array([channel.alpha_v for channel in channels.values()])
        u_v = np.array([channel.u_v for channel in channels.values()])
        v_v = np.array([channel.v_v for channel in channels.values()])
        T_v = np.array([channel.T_v for channel in channels.values()])
        
        # 1. 气相分数分布
        ax1 = fig.add_subplot(221, projection='3d')
        self.plot_3d_field(ax1, X, Y, alpha_v,
                          '气相分数分布', '轴向位置 (m)', '子通道编号', '气相分数')
        
        # 2. 气相轴向速度分布
        ax2 = fig.add_subplot(222, projection='3d')
        self.plot_3d_field(ax2, X, Y, u_v,
                          '气相轴向速度分布', '轴向位置 (m)', '子通道编号', '速度 (m/s)')
        
        # 3. 气相横向速度分布
        ax3 = fig.add_subplot(223, projection='3d')
        self.plot_3d_field(ax3, X, Y, v_v,
                          '气相横向速度分布', '轴向位置 (m)', '子通道编号', '速度 (m/s)')
        
        # 4. 气相��度分布
        ax4 = fig.add_subplot(224, projection='3d')
        self.plot_3d_field(ax4, X, Y, T_v,
                          '气相温度分布', '轴向位置 (m)', '子通道编号', '温度 (K)')
        
        plt.tight_layout()
        plt.show()
        
    def plot_axial_profiles(self, channels, channel_id):
        """
        绘制单个子通道的轴向分布
        
        参数：
        channels: dict, 子通道字典
        channel_id: int, 子通道编号
        """
        channel = channels[channel_id]
        z = np.linspace(0, cfg.CHANNEL_LENGTH, cfg.AXIAL_NODES)
        
        # 创建图形窗口
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # 1. 相分数分布
        ax1.plot(z, channel.alpha_l, 'b-', label='液相')
        ax1.plot(z, channel.alpha_v, 'r--', label='气相')
        ax1.set_title(f'子通道 {channel_id} 的相分数分布')
        ax1.set_xlabel('轴向位置 (m)')
        ax1.set_ylabel('相分数')
        ax1.grid(True)
        ax1.legend()
        
        # 2. 速度分布
        ax2.plot(z, channel.u_l, 'b-', label='液相轴向')
        ax2.plot(z, channel.u_v, 'r--', label='气相轴向')
        ax2.plot(z, channel.v_l, 'g-.', label='液相横向')
        ax2.plot(z, channel.v_v, 'm:', label='气相横向')
        ax2.set_title(f'子通道 {channel_id} 的速度分布')
        ax2.set_xlabel('轴向位置 (m)')
        ax2.set_ylabel('速度 (m/s)')
        ax2.grid(True)
        ax2.legend()
        
        # 3. 温度分布
        ax3.plot(z, channel.T_l, 'b-', label='液相')
        ax3.plot(z, channel.T_v, 'r--', label='气相')
        ax3.plot(z, channel.T_w, 'k-.', label='壁面')
        ax3.set_title(f'子通道 {channel_id} 的温度分布')
        ax3.set_xlabel('轴向位置 (m)')
        ax3.set_ylabel('温度 (K)')
        ax3.grid(True)
        ax3.legend()
        
        # 4. 压力分布
        ax4.plot(z, channel.pressure/1e6, 'k-')
        ax4.set_title(f'子通道 {channel_id} 的压力分布')
        ax4.set_xlabel('轴向位置 (m)')
        ax4.set_ylabel('压力 (MPa)')
        ax4.grid(True)
        
        plt.tight_layout()
        plt.show()
        
    def plot_channel_comparison(self, channels, variable='T_l'):
        """
        比较不同子通道的某个变量分布
        
        参数：
        channels: dict, 子通道字典
        variable: str, 要比较的变量名
        """
        z = np.linspace(0, cfg.CHANNEL_LENGTH, cfg.AXIAL_NODES)
        
        plt.figure(figsize=(10, 6))
        for channel_id, channel in channels.items():
            values = getattr(channel, variable)
            plt.plot(z, values, label=f'通道 {channel_id}')
            
        plt.title(f'各子通道 {variable} 的对比')
        plt.xlabel('轴向位置 (m)')
        plt.ylabel(variable)
        plt.grid(True)
        plt.legend()
        plt.show()
        
    def save_results(self, channels, filename):
        """
        保存计算结果为图片
        
        参数：
        channels: dict, 子通道字典
        filename: str, 文件名
        """
        # 绘制液相结果
        self.plot_liquid_phase_results(channels)
        plt.savefig(f'{filename}_liquid.png')
        plt.close()
        
        # 绘制气相结果
        self.plot_vapor_phase_results(channels)
        plt.savefig(f'{filename}_vapor.png')
        plt.close()
        
        # 绘制每个通道的轴向分布
        for channel_id in channels.keys():
            self.plot_axial_profiles(channels, channel_id)
            plt.savefig(f'{filename}_channel_{channel_id}.png')
            plt.close() 